"""
WebSocket Connection Manager
Handles WebSocket connections for real-time communication during PPT generation.
"""

import json
import logging
from typing import Dict, List, Optional, Any
from fastapi import WebSocket, WebSocketDisconnect
from datetime import datetime
import asyncio

logger = logging.getLogger(__name__)


class ConnectionManager:
    """
    Manages WebSocket connections for real-time communication
    
    Handles connection lifecycle, message broadcasting, and session management
    """
    
    def __init__(self):
        # Store active connections by session_id
        self.active_connections: Dict[str, WebSocket] = {}
        # Store connection metadata
        self.connection_info: Dict[str, Dict[str, Any]] = {}
        
    async def connect(self, websocket: WebSocket, session_id: str) -> bool:
        """
        Accept a new WebSocket connection and register it
        
        Args:
            websocket: WebSocket connection instance
            session_id: Unique session identifier
            
        Returns:
            bool: True if connection successful, False otherwise
        """
        try:
            await websocket.accept()
            
            # Close existing connection for this session if any
            if session_id in self.active_connections:
                await self.disconnect(session_id)
            
            # Register new connection
            self.active_connections[session_id] = websocket
            self.connection_info[session_id] = {
                "connected_at": datetime.now(),
                "last_ping": datetime.now(),
                "message_count": 0
            }
            
            logger.info(f"WebSocket connected for session: {session_id}")
            
            # Send connection confirmation
            await self.send_personal_message(session_id, {
                "type": "connection_established",
                "session_id": session_id,
                "timestamp": datetime.now().isoformat()
            })
            
            return True
            
        except Exception as e:
            logger.error(f"Failed to connect WebSocket for session {session_id}: {e}")
            return False
    
    async def disconnect(self, session_id: str):
        """
        Disconnect and cleanup a WebSocket connection
        
        Args:
            session_id: Session identifier to disconnect
        """
        if session_id in self.active_connections:
            try:
                websocket = self.active_connections[session_id]
                await websocket.close()
            except Exception as e:
                logger.warning(f"Error closing WebSocket for session {session_id}: {e}")
            finally:
                # Clean up connection data
                del self.active_connections[session_id]
                if session_id in self.connection_info:
                    del self.connection_info[session_id]
                
                logger.info(f"WebSocket disconnected for session: {session_id}")
    
    async def send_personal_message(self, session_id: str, message: Dict[str, Any]) -> bool:
        """
        Send a message to a specific session
        
        Args:
            session_id: Target session identifier
            message: Message data to send
            
        Returns:
            bool: True if message sent successfully, False otherwise
        """
        if session_id not in self.active_connections:
            logger.warning(f"No active connection for session: {session_id}")
            return False
        
        try:
            websocket = self.active_connections[session_id]
            
            # Add metadata to message
            message_with_meta = {
                **message,
                "timestamp": datetime.now().isoformat(),
                "session_id": session_id
            }
            
            await websocket.send_text(json.dumps(message_with_meta))
            
            # Update connection stats
            if session_id in self.connection_info:
                self.connection_info[session_id]["message_count"] += 1
                self.connection_info[session_id]["last_ping"] = datetime.now()
            
            return True
            
        except WebSocketDisconnect:
            logger.info(f"WebSocket disconnected during message send: {session_id}")
            await self.disconnect(session_id)
            return False
        except Exception as e:
            logger.error(f"Failed to send message to session {session_id}: {e}")
            return False
    
    async def broadcast_message(self, message: Dict[str, Any], exclude_sessions: Optional[List[str]] = None):
        """
        Broadcast a message to all active connections
        
        Args:
            message: Message data to broadcast
            exclude_sessions: List of session IDs to exclude from broadcast
        """
        if not self.active_connections:
            return
        
        exclude_sessions = exclude_sessions or []
        message_with_meta = {
            **message,
            "timestamp": datetime.now().isoformat(),
            "type": "broadcast"
        }
        
        # Send to all active connections
        disconnected_sessions = []
        for session_id, websocket in self.active_connections.items():
            if session_id in exclude_sessions:
                continue
                
            try:
                await websocket.send_text(json.dumps(message_with_meta))
                
                # Update connection stats
                if session_id in self.connection_info:
                    self.connection_info[session_id]["message_count"] += 1
                    
            except WebSocketDisconnect:
                disconnected_sessions.append(session_id)
            except Exception as e:
                logger.error(f"Failed to broadcast to session {session_id}: {e}")
                disconnected_sessions.append(session_id)
        
        # Clean up disconnected sessions
        for session_id in disconnected_sessions:
            await self.disconnect(session_id)
    
    async def send_progress_update(self, session_id: str, progress: int, step: str, status: str):
        """
        Send a progress update message to a specific session
        
        Args:
            session_id: Target session identifier
            progress: Progress percentage (0-100)
            step: Current step description
            status: Current status
        """
        message = {
            "type": "progress_update",
            "progress": progress,
            "step": step,
            "status": status
        }
        await self.send_personal_message(session_id, message)
    
    async def send_error_message(self, session_id: str, error: str, error_code: Optional[str] = None):
        """
        Send an error message to a specific session
        
        Args:
            session_id: Target session identifier
            error: Error message
            error_code: Optional error code
        """
        message = {
            "type": "error",
            "error": error,
            "error_code": error_code
        }
        await self.send_personal_message(session_id, message)
    
    async def send_completion_message(self, session_id: str, file_path: str, slides_count: int):
        """
        Send a completion message to a specific session
        
        Args:
            session_id: Target session identifier
            file_path: Path to generated PPT file
            slides_count: Number of slides generated
        """
        message = {
            "type": "generation_complete",
            "file_path": file_path,
            "slides_count": slides_count,
            "status": "completed"
        }
        await self.send_personal_message(session_id, message)
    
    def get_connection_count(self) -> int:
        """Get the number of active connections"""
        return len(self.active_connections)
    
    def get_connection_info(self, session_id: str) -> Optional[Dict[str, Any]]:
        """Get connection information for a specific session"""
        return self.connection_info.get(session_id)
    
    def is_connected(self, session_id: str) -> bool:
        """Check if a session has an active connection"""
        return session_id in self.active_connections
    
    async def ping_all_connections(self):
        """Send ping to all connections to check if they're still alive"""
        if not self.active_connections:
            return
        
        ping_message = {
            "type": "ping",
            "timestamp": datetime.now().isoformat()
        }
        
        disconnected_sessions = []
        for session_id, websocket in self.active_connections.items():
            try:
                await websocket.send_text(json.dumps(ping_message))
            except (WebSocketDisconnect, Exception):
                disconnected_sessions.append(session_id)
        
        # Clean up disconnected sessions
        for session_id in disconnected_sessions:
            await self.disconnect(session_id)


# Global connection manager instance
connection_manager = ConnectionManager()